import cv2
import numpy as np
from noise import Noise


def MAP_Sobolev(g, K):
    G = np.fft.fft2(g)

    h, w = g.shape

    C = np.arange(0, h).reshape((-1, 1)) ** 2 + np.arange(0, w) ** 2

    F_fit = G / (1 + K * C)

    f_fit = np.fft.ifft2(F_fit)

    # f_fit = np.fft.fftshift(f_fit)

    f_fit = np.abs(f_fit)

    f_fit = ((f_fit - f_fit.min()) / (f_fit.max() - f_fit.min()) * 255).astype('uint8')

    return f_fit


def MAP_TV(g, lr, n, k):

    f_fit = g

    for _ in range(n):
        x = cv2.Sobel(f_fit, cv2.CV_32F, 1, 0)
        y = cv2.Sobel(f_fit, cv2.CV_32F, 0, 1)
        l = np.sqrt(x ** 2 + y ** 2 + 1e-6)
        x /= l
        y /= l
        x = cv2.Sobel(x, cv2.CV_32F, 1, 0)
        y = cv2.Sobel(y, cv2.CV_32F, 1, 0)
        div = x + y

        f_fit = f_fit - lr * (2 * (g - f_fit) - k * div)
    
    f_fit = ((f_fit - f_fit.min()) / (f_fit.max() - f_fit.min()) * 255).astype('uint8')

    return f_fit


def MLE():
    pass


if __name__ == '__main__':
    SRC = cv2.imread('peppers-bw.bmp', cv2.IMREAD_GRAYSCALE)

    def loss(y, label):
        return 10 * np.log10(255 ** 2 / np.mean((label - y) ** 2 + 1e-6))

    def show(winname, img):
        cv2.imshow(winname, img)
        cv2.resizeWindow(winname, 512, 512)
        cv2.waitKey(1)
        print(winname, str(loss(SRC, img)))

    # 噪声模型
    noise = Noise()
    # 参数
    sigma = 5
    # 退化图像
    g = noise(SRC, sigma=sigma)
    
    G = np.fft.fft2(g)

    show('SRC', SRC)

    show('NOISE', g)

    show('MAP_Sobolev', MAP_Sobolev(g, 1e-5))

    show('MAP_TV', MAP_TV(g, 1e-2, 100, 1e-5))

    cv2.waitKey(0)